feat: add ai/embed connector for RAG#167
Conversation
Adds a text-embeddings connector as the foundation for retrieval-augmented generation. - ai/embed connector (internal/connector/embed.go): provider selection, base-URL/model allowlists, and AI metrics mirroring ai/completion. It lives under the "ai/" prefix so the engine's token-budget accounting applies automatically (it returns usage.total_tokens). - EmbeddingProvider interface + OpenAI implementation (also covers Azure OpenAI and OpenAI-compatible servers via base_url). Extracted the shared OpenAI credential-header logic used by chat and embeddings. - Output includes a pgvector text literal (output.vector, "[...]") so an embedding drops straight into a postgres/query arg cast with ::vector, removing the array-encoding friction of building RAG by hand. - Registered ai/embed; serve applies the base-URL allowlist to it. - Example: pgvector RAG (rag-kb-schema.sql, rag-ingest.yaml, rag-ask.yaml), a rag-guide doc, and an ai/embed connector-reference entry. Both workflows pass `mantle validate`. Scoped as the first RAG increment. Bedrock embeddings (InvokeModel) and native kb/upsert / kb/query + chunking are planned follow-ups on #153. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Warning Review limit reached
Next review available in: 29 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (12)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 07a6e7cdb8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| out := make([][]float64, len(apiResp.Data)) | ||
| for _, d := range apiResp.Data { | ||
| if d.Index < 0 || d.Index >= len(out) { | ||
| return nil, fmt.Errorf("openai: embedding index %d out of range", d.Index) | ||
| } |
There was a problem hiding this comment.
When an OpenAI-compatible endpoint returns fewer data items than requested or repeats an index, this builds the result slice from len(apiResp.Data) and never verifies that each requested input got exactly one embedding. A multi-input ai/embed step can then silently return a shorter vectors list or include an empty [] vector for a missing slot, causing downstream pgvector writes/queries to be associated with the wrong input instead of failing fast; size the output by len(req.Inputs) and reject missing/duplicate indexes.
Useful? React with 👍 / 👎.
Addresses PR review (Codex P2). The OpenAI embeddings parser sized the result by len(data) and indexed by the API's index field, so a response with fewer items than inputs — or a duplicate index — could silently misalign embeddings with their inputs or leave an empty [] vector, corrupting downstream pgvector writes/queries. Size the output by len(req.Inputs) and reject out-of-range, duplicate, or missing indexes. Adds a regression test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
2 participants
Summary
First increment on RAG (#153): an
ai/embedtext-embeddings connector, plus a working pgvector RAG example. Together with the existingai/completion, this lets you build retrieval-augmented generation on stock connectors — embed → store → nearest-neighbour search → grounded answer.Changes
ai/embedconnector (internal/connector/embed.go) — provider selection, base-URL/model allowlists, and AI metrics mirroringai/completion. It lives under theai/prefix, so the engine's token-budget accounting applies automatically (it returnsusage.total_tokens).EmbeddingProviderinterface + OpenAI implementation — covers OpenAI, Azure OpenAI, and OpenAI-compatible servers viabase_url. Extracted the shared OpenAI credential-header logic so chat and embeddings reuse it.output.embedding) and as a pgvector text literal (output.vector,"[0.1,0.2,...]"), which binds straight into apostgres/queryarg cast with::vector. This removes the array-encoding friction of hand-building RAG.ai/embed;serveapplies the base-URL allowlist to it (the model allowlist is intentionally not shared, since it lists chat models).rag-kb-schema.sql(pgvector table + HNSW index),rag-ingest.yaml,rag-ask.yaml, arag-guidedoc, and anai/embedentry in the connector reference.Scope
This is the foundational slice of #153. Deliberately deferred (follow-ups on #153): Bedrock embeddings (needs
InvokeModel+ per-model shapes, distinct from the Converse API used for chat), and nativekb/upsert/kb/queryconnectors + chunking. RAG is fully composable today fromai/embed+postgres/query+ai/completion, as the example shows.Testing
go test ./internal/connector/ -run TestEmbedding(httptest-based) passes, plus the existing OpenAI provider tests (credential refactor is covered).go build,go vet,gofmt, andgolangci-lint runclean.rag-ingest.yamlandrag-ask.yamlpassmantle validate.Related Issues
Part of #153 (RAG subsystem).
Generated by Claude Code